package com.huatai.datacenter.utils;

import com.alibaba.fastjson.JSONObject;
import com.huatai.datacenter.entity.BrokerEntity;
import kafka.utils.ZKStringSerializer;
import lombok.extern.slf4j.Slf4j;
import org.I0Itec.zkclient.ZkClient;
import org.I0Itec.zkclient.exception.ZkMarshallingError;
import org.I0Itec.zkclient.serialize.ZkSerializer;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.*;

/**
 * @author Lion
 * @date 2023/3/20  13:39
 */
@Slf4j
public class ZKUtil implements Cloneable {

	private ZkClient zkClient;

	private int SESSION_TIMEOUT = 1000 * 60;

	private int CONNECTION_TIMEOUT = 1000 * 60;

	private static final String BROKER_PATH = "/brokers/topics/";

	private static final String TOPIC_CONFIG_CHANGES_PATH = "/config/changes";

	private static final String TOPIC_CONFIG_CHANGE_ZNODE_PREFIX = "config_change_";

	private String rootPath = "";

	public ZKUtil(String zkServers){

		String lastChar = zkServers.substring(zkServers.length() -1,zkServers.length());
		//如果最后一个字符是‘/’，去掉
		if ("/".equals(lastChar)){
			zkServers = zkServers.substring(0,zkServers.length()-1);
		}
		//获取指定path
		int index = zkServers.indexOf("/");
		if (index != -1){
			//存在rootpath
			String newzkServers = zkServers.substring(0,index);
			this.rootPath = zkServers.substring(index);
			zkClient = new ZkClient(newzkServers,SESSION_TIMEOUT,CONNECTION_TIMEOUT,new ZkStringSerialize());
		}else {
			zkClient = new ZkClient(zkServers,SESSION_TIMEOUT,CONNECTION_TIMEOUT,new ZkStringSerialize());
		}
	}

	public ZKUtil(String zkServers,int sessionTimeout,int connectionTimeOut){
		String lastChar = zkServers.substring(zkServers.length()-1, zkServers.length());
		if("/".equals(lastChar)) {
			zkServers = zkServers.substring(0,zkServers.length()-1);
		}
		int index = zkServers.indexOf("/");
		if(index != -1) {
			// 存在rootPath
			String newzkServers = zkServers.substring(0,index);
			this.rootPath = zkServers.substring(index);
			zkClient = new ZkClient(newzkServers, sessionTimeout, connectionTimeOut, new ZkStringSerialize());
		} else {
			zkClient = new ZkClient(zkServers, sessionTimeout, connectionTimeOut, new ZkStringSerialize());
		}
	}

	/**
	 * @author Lion
	 * @description 根据group与topic获取相应的partition，offset提交情况
	 * @param groupId
	 * @param topic
	 * @return java.util.Map<java.lang.String,java.util.Map<java.lang.String,java.lang.String>>
	 */
	public Map<String,Map<String,String>> getZKConsumerOffsets(String groupId,String topic){
		Map<String,Map<String,String>> result = new HashMap<>();
		String offsetsPath = rootPath + "/consumers/" + groupId + "/offsets/" + topic;

		if (zkClient.exists(offsetsPath)){
			List<String> offsets = zkClient.getChildren(offsetsPath);
			offsets.forEach(partition -> {

				try {
					String offset = zkClient.readData(offsetsPath + "/" + partition,true);
					if (offset != null){
						Map<String,String> map = new HashMap<>();
						map.put("offset",offset);
						result.put(partition,map);
					}
				}catch (Exception e){
					log.error("",e);
				}
			});
		}

		String ownersPath = rootPath + "/consumers/" + groupId + "/owners/" + topic;
		if (zkClient.exists(ownersPath)){
			List<String> owners = zkClient.getChildren(ownersPath);
			if (owners != null && !owners.isEmpty()){
				owners.forEach(partition -> {
					try {
						String owner = zkClient.readData(ownersPath + "/" + partition,true);
						if (owner != null){
							Map<String,String> map = result.get(partition);
							if (map != null){
								map.put("owner",owner);
								result.put(partition,map);
							}
						}
					}catch (Exception e){
						log.error("",e);
					}
				});
			}
		}
		return result;
	}

	/**
	 * @author Lion
	 * @description 获取zookeeper上消费offset信息
	 * @param  * @param
	 * @return java.util.Map<java.lang.String,java.util.Map<java.lang.String,java.util.Map<java.lang.String,java.lang.String>>>
	 */
	public Map<String,Map<String,Map<String,String>>> getZKConsumerOffsets(){
		Map<String,Map<String,Map<String,String>>> result = new HashMap<>();
		List<String> allGroups = zkClient.getChildren(rootPath + "/consumer");

		allGroups.forEach(group ->{
			if (zkClient.exists(rootPath + "/consumer/" + group + "offset")){
				Map<String,Map<String,String>> topicOffsets = new HashMap<>();
				Set<String> topics = new HashSet<>(zkClient.getChildren(rootPath + "/consumers" + group + "/offsets/"));
				topics.forEach(topic ->{
					List<String> partitions = zkClient.getChildren(rootPath + "/consumers" + group + "/offsets" + topic);
					Map<String,String> offsets = new HashMap<>();
					Map<String,String> results = getState(topic,group);
					partitions.forEach(partition -> {
						String path = rootPath + "/consumers/" + group + "offsets" + topic + "/" + partition;
						String offset = zkClient.readData(path,true);
						String state = results.get(partition);
						offsets.put(partition,offset + "|" + state);
					});
					topicOffsets.put(topic,offsets);
				});
				result.put(group,topicOffsets);
			}
		});

		return result;
	}

	/**
	 * @author Lion
	 * @description 获取状态
	 * @param topic
	 * @param groupId
	 * @return java.util.Map<java.lang.String,java.lang.String>
	 */
	public Map<String,String> getState(String topic,String groupId){
		Map<String,String> result = new HashMap<>();
		String ownersPath = rootPath + "/consumer/" + groupId + "/owners/" + topic;

		if (zkClient.exists(ownersPath)){
			List<String> owners = zkClient.getChildren(ownersPath);
			for (String partition:owners){
				String owner = zkClient.readData(ownersPath + "/" + partition,true);
				if (owner !=null){
					result.put(partition,owner);
				}
			}
		}
		return result;
	}

	/**
	 * @author Lion
	 * @description 获取所有消费组
	 * @param
	 * @return java.util.Set<java.lang.String>
	 */
	public Set<String> listConsumerGroups(){
		List<String> allGroups = new ArrayList<>();
		try {
			allGroups = zkClient.getChildren(rootPath + "/consumers");
		}catch (Exception e){
			return new HashSet<>(allGroups);
		}

		return new HashSet<>(allGroups);
	}

	/**
	 * @author Lion
	 * @description 根据主题获取消费组
	 * @param  topic
	 * @return java.util.Set<java.lang.String>
	 */
	public Set<String> listTopicGroups(String topic){
		Set<String> groups = new HashSet<>();

		List<String> allGroups = zkClient.getChildren(rootPath + "/consumers");
		allGroups.forEach(group ->{
			if (zkClient.exists(rootPath + "/consumers/" + group + "/offsets")){
				Set<String> offsets = new HashSet<>(zkClient.getChildren(rootPath + "/consumers/" + group + "/offsets"));
				if (offsets.contains(topic)){
					groups.add(group);
				}
			}
		});
		return groups;
	}

	/**
	 * @author Lion
	 * @description 根据消费组获取topic
	 * @param group 消费组
	 * @return java.util.Set<java.lang.String>
	 */
	public Set<String> listTopicByGroup(String group){
		Set<String> topics = new HashSet<>();

		topics = new HashSet<>(zkClient.getChildren(rootPath + "/consumers/" + group + "offsets"));

		return new HashSet<>(topics);
	}

	/**
	 * @author Lion
	 * @description 判断group能否删除
	 * @param consumerGroup
	 * @return boolean
	 */
	private boolean groupCanDelete(String consumerGroup){
		Set<String> topics = listTopicByGroup(consumerGroup);
		String ownersPath;

		for (String topic:topics){
			ownersPath = rootPath + "/consumers/" + consumerGroup + "/owners/" + topic;
			if (zkClient.exists(ownersPath)){
				List<String> partitions = zkClient.getChildren(ownersPath);
				for (String partiton : partitions){
					String owner = zkClient.readData(ownersPath + "/" + partiton,true);
					if (null != owner){
						return false;
					}
				}
			}
		}
		return true;
	}

	public  List<BrokerEntity> getBrokers() throws Exception {
		List<BrokerEntity> kafkaHosts = new ArrayList<BrokerEntity>();
		try {
			List<String> ids = zkClient.getChildren(rootPath + "/brokers/ids");
			for (String id : ids) {
				String brokerInfo = readDataMaybeNull(rootPath + "/brokers/ids/" + id).getData();
				BrokerEntity bi = new BrokerEntity();
				bi.setBid(Integer.parseInt(id));
				JSONObject jsonObj = JSONObject.parseObject(brokerInfo);
				if (jsonObj.containsKey("host")) {
					bi.setHost(jsonObj.get("host").toString());
				}
				if (jsonObj.containsKey("port")) {
					bi.setPort(Integer.parseInt(jsonObj.get("port").toString()));
				}
				if (jsonObj.containsKey("jmx_port")) {
					bi.setJmxPort(Integer.parseInt(jsonObj.get("jmx_port").toString()));
				}
				if (jsonObj.containsKey("version")) {
					bi.setVersion(Integer.parseInt(jsonObj.get("version").toString()));
				}
				if (jsonObj.containsKey("timestamp")) {
					bi.setTimestamp(Long.parseLong(jsonObj.get("timestamp").toString()));
				}
				kafkaHosts.add(bi);
			}
		}catch (Exception e) {
		}
		return kafkaHosts;
	}


	class ZkStringSerialize implements ZkSerializer {
		@Override
		public byte[] serialize(Object paramObject) throws ZkMarshallingError {
			return ZKStringSerializer.serialize(paramObject);
		}

		@Override
		public Object deserialize(byte[] paramArrayOfByte) throws ZkMarshallingError {
			return ZKStringSerializer.deserialize(paramArrayOfByte);
		}
	}

	public  ZkDataAndStat readDataMaybeNull(String path) {
		Stat stat = new Stat();
		String data = null;
		try {
			data = zkClient.readData(path, stat);
		} catch (Exception e) {
			log.warn("zk read data error." ,e);
		}
		return new ZkDataAndStat(data, stat);
	}

	class ZkDataAndStat {
		private String data;
		private Stat stat;
		public String getData() {
			return data;
		}
		public void setData(String data) {
			this.data = data;
		}
		public Stat getStat() {
			return stat;
		}
		public void setStat(Stat stat) {
			this.stat = stat;
		}
		public ZkDataAndStat(String data, Stat stat) {
			super();
			this.data = data;
			this.stat = stat;
		}
	}

	public void close() throws IOException {
		if (zkClient != null) {
			zkClient.close();
		}
	}

}
